"""
# @Time    : 2021/7/1 7:15 下午
# @Author  : hezhiqiang01
# @Email   : hezhiqiang01@baidu.com
# @File    : env_runner.py
"""

import time
import numpy as np
import torch
from runner.shared.base_runner import Runner

# import imageio


def _t2n(x):
    """
    将PyTorch张量转换为NumPy数组
    
    参数:
        x: PyTorch张量
        
    返回:
        NumPy数组
    """
    return x.detach().cpu().numpy()


class EnvRunner(Runner):
    """
    用于执行训练、评估和数据收集的运行器类。详细信息请参见父类。
    """

    def __init__(self, config):
        """
        初始化环境运行器
        
        参数:
            config: 配置字典
        """
        super(EnvRunner, self).__init__(config)

    def run(self):
        """
        运行训练过程
        """
        self.warmup()  # 预热阶段

        start = time.time()  # 记录开始时间
        episodes = int(self.num_env_steps) // self.episode_length // self.n_rollout_threads  # 计算总回合数

        for episode in range(episodes):  # 遍历每个回合
            if self.use_linear_lr_decay:  # 如果使用线性学习率衰减
                self.trainer.policy.lr_decay(episode, episodes)  # 学习率衰减

            for step in range(self.episode_length):  # 遍历每个步骤
                # 采样动作
                (
                    values,
                    actions,
                    action_log_probs,
                    rnn_states,
                    rnn_states_critic,
                    actions_env,
                ) = self.collect(step)  # 收集数据

                # 观察奖励和下一个状态
                obs, rewards, dones, infos = self.envs.step(actions_env)  # 执行环境步进

                data = (
                    obs,
                    rewards,
                    dones,
                    infos,
                    values,
                    actions,
                    action_log_probs,
                    rnn_states,
                    rnn_states_critic,
                )  # 打包数据

                # 将数据插入缓冲区
                self.insert(data)  # 插入数据

            # 计算回报并更新网络
            self.compute()  # 计算回报
            train_infos = self.train()  # 训练网络

            # 后处理
            total_num_steps = (episode + 1) * self.episode_length * self.n_rollout_threads  # 计算总步数

            # 保存模型
            if episode % self.save_interval == 0 or episode == episodes - 1:  # 如果达到保存间隔或最后一个回合
                self.save()  # 保存模型

            # 记录信息
            if episode % self.log_interval == 0:  # 如果达到日志间隔
                end = time.time()  # 记录结束时间
                print(
                    "\n 场景 {} 算法 {} 实验 {} 更新 {}/{} 回合, 总时间步数 {}/{}, FPS {}.\n".format(
                        self.all_args.scenario_name,
                        self.algorithm_name,
                        self.experiment_name,
                        episode,
                        episodes,
                        total_num_steps,
                        self.num_env_steps,
                        int(total_num_steps / (end - start)),
                    )
                )  # 打印信息

                # if self.env_name == "MPE":
                #     env_infos = {}
                #     for agent_id in range(self.num_agents):
                #         idv_rews = []
                #         for info in infos:
                #             if 'individual_reward' in info[agent_id].keys():
                #                 idv_rews.append(info[agent_id]['individual_reward'])
                #         agent_k = 'agent%i/individual_rewards' % agent_id
                #         env_infos[agent_k] = idv_rews

                train_infos["average_episode_rewards"] = np.mean(self.buffer.rewards) * self.episode_length  # 计算平均回合奖励
                print("平均回合奖励为 {}".format(train_infos["average_episode_rewards"]))  # 打印平均回合奖励
                self.log_train(train_infos, total_num_steps)  # 记录训练信息
                # self.log_env(env_infos, total_num_steps)

            # 评估
            if episode % self.eval_interval == 0 and self.use_eval:  # 如果达到评估间隔且使用评估
                self.eval(total_num_steps)  # 评估

    def warmup(self):
        """
        预热阶段，收集初始数据
        """
        # 重置环境
        obs = self.envs.reset()  # shape = [env_num, agent_num, obs_dim]

        # 回放缓冲区
        if self.use_centralized_V:  # 如果使用中心化价值函数
            share_obs = obs.reshape(self.n_rollout_threads, -1)  # shape = [env_num, agent_num * obs_dim]
            share_obs = np.expand_dims(share_obs, 1).repeat(
                self.num_agents, axis=1
            )  # shape = shape = [env_num, agent_num， agent_num * obs_dim]
        else:
            share_obs = obs  # 否则共享观察就是观察

        self.buffer.share_obs[0] = share_obs.copy()  # 存储共享观察
        self.buffer.obs[0] = obs.copy()  # 存储观察

    @torch.no_grad()
    def collect(self, step):
        """
        收集数据
        
        参数:
            step: 当前步数
            
        返回:
            values: 价值
            actions: 动作
            action_log_probs: 动作对数概率
            rnn_states: RNN状态
            rnn_states_critic: 评论家RNN状态
            actions_env: 环境动作
        """
        self.trainer.prep_rollout()  # 准备收集数据
        (
            value,
            action,
            action_log_prob,
            rnn_states,
            rnn_states_critic,
        ) = self.trainer.policy.get_actions(
            np.concatenate(self.buffer.share_obs[step]),  # 连接共享观察
            np.concatenate(self.buffer.obs[step]),  # 连接观察
            np.concatenate(self.buffer.rnn_states[step]),  # 连接RNN状态
            np.concatenate(self.buffer.rnn_states_critic[step]),  # 连接评论家RNN状态
            np.concatenate(self.buffer.masks[step]),  # 连接掩码
        )  # 获取动作
        # [self.envs, agents, dim]
        values = np.array(np.split(_t2n(value), self.n_rollout_threads))  # [env_num, agent_num, 1]
        actions = np.array(np.split(_t2n(action), self.n_rollout_threads))  # [env_num, agent_num, action_dim]
        action_log_probs = np.array(
            np.split(_t2n(action_log_prob), self.n_rollout_threads)
        )  # [env_num, agent_num, 1]
        rnn_states = np.array(
            np.split(_t2n(rnn_states), self.n_rollout_threads)
        )  # [env_num, agent_num, 1, hidden_size]
        rnn_states_critic = np.array(
            np.split(_t2n(rnn_states_critic), self.n_rollout_threads)
        )  # [env_num, agent_num, 1, hidden_size]
        # 重新排列动作
        if self.envs.action_space[0].__class__.__name__ == "MultiDiscrete":  # 如果是多离散动作空间
            for i in range(self.envs.action_space[0].shape):
                uc_actions_env = np.eye(self.envs.action_space[0].high[i] + 1)[actions[:, :, i]]  # one-hot编码
                if i == 0:
                    actions_env = uc_actions_env  # 第一个动作
                else:
                    actions_env = np.concatenate((actions_env, uc_actions_env), axis=2)  # 连接动作
        elif self.envs.action_space[0].__class__.__name__ == "Discrete":  # 如果是离散动作空间
            # actions  --> actions_env : shape:[10, 1] --> [5, 2, 5]
            actions_env = np.squeeze(np.eye(self.envs.action_space[0].n)[actions], 2)  # one-hot编码
        else:
            # TODO 这里改造成自己环境需要的形式即可
            # TODO 在这里，你可以将actions_env的形状更改为适合你的环境
            actions_env = actions  # 直接使用动作
            # raise NotImplementedError

        return (
            values,
            actions,
            action_log_probs,
            rnn_states,
            rnn_states_critic,
            actions_env,
        )  # 返回结果

    def insert(self, data):
        """
        插入数据到缓冲区
        
        参数:
            data: 要插入的数据
        """
        (
            obs,
            rewards,
            dones,
            infos,
            values,
            actions,
            action_log_probs,
            rnn_states,
            rnn_states_critic,
        ) = data  # 解包数据

        rnn_states[dones == True] = np.zeros(
            ((dones == True).sum(), self.recurrent_N, self.hidden_size),
            dtype=np.float32,
        )  # 重置已完成环境的RNN状态
        rnn_states_critic[dones == True] = np.zeros(
            ((dones == True).sum(), *self.buffer.rnn_states_critic.shape[3:]),
            dtype=np.float32,
        )  # 重置已完成环境的评论家RNN状态
        masks = np.ones((self.n_rollout_threads, self.num_agents, 1), dtype=np.float32)  # 创建掩码
        masks[dones == True] = np.zeros(((dones == True).sum(), 1), dtype=np.float32)  # 设置已完成环境的掩码为0

        if self.use_centralized_V:  # 如果使用中心化价值函数
            share_obs = obs.reshape(self.n_rollout_threads, -1)  # 重塑观察
            share_obs = np.expand_dims(share_obs, 1).repeat(self.num_agents, axis=1)  # 扩展观察
        else:
            share_obs = obs  # 否则共享观察就是观察

        self.buffer.insert(
            share_obs,
            obs,
            rnn_states,
            rnn_states_critic,
            actions,
            action_log_probs,
            values,
            rewards,
            masks,
        )  # 插入数据到缓冲区

    @torch.no_grad()
    def eval(self, total_num_steps):
        """
        评估智能体
        
        参数:
            total_num_steps: 总步数
        """
        eval_episode_rewards = []  # 评估回合奖励
        eval_obs = self.eval_envs.reset()  # 重置评估环境

        eval_rnn_states = np.zeros(
            (self.n_eval_rollout_threads, *self.buffer.rnn_states.shape[2:]),
            dtype=np.float32,
        )  # 初始化评估RNN状态
        eval_masks = np.ones((self.n_eval_rollout_threads, self.num_agents, 1), dtype=np.float32)  # 初始化评估掩码

        for eval_step in range(self.episode_length):  # 遍历评估步数
            self.trainer.prep_rollout()  # 准备收集数据
            eval_action, eval_rnn_states = self.trainer.policy.act(
                np.concatenate(eval_obs),  # 连接观察
                np.concatenate(eval_rnn_states),  # 连接RNN状态
                np.concatenate(eval_masks),  # 连接掩码
                deterministic=True,  # 确定性动作
            )  # 获取动作
            eval_actions = np.array(np.split(_t2n(eval_action), self.n_eval_rollout_threads))  # 分割动作
            eval_rnn_states = np.array(np.split(_t2n(eval_rnn_states), self.n_eval_rollout_threads))  # 分割RNN状态

            if self.eval_envs.action_space[0].__class__.__name__ == "MultiDiscrete":  # 如果是多离散动作空间
                for i in range(self.eval_envs.action_space[0].shape):
                    eval_uc_actions_env = np.eye(self.eval_envs.action_space[0].high[i] + 1)[
                        eval_actions[:, :, i]
                    ]  # one-hot编码
                    if i == 0:
                        eval_actions_env = eval_uc_actions_env  # 第一个动作
                    else:
                        eval_actions_env = np.concatenate((eval_actions_env, eval_uc_actions_env), axis=2)  # 连接动作
            elif self.eval_envs.action_space[0].__class__.__name__ == "Discrete":  # 如果是离散动作空间
                eval_actions_env = np.squeeze(np.eye(self.eval_envs.action_space[0].n)[eval_actions], 2)  # one-hot编码
            else:
                raise NotImplementedError  # 不支持的动作空间

            # 观察奖励和下一个状态
            eval_obs, eval_rewards, eval_dones, eval_infos = self.eval_envs.step(eval_actions_env)  # 执行环境步进
            eval_episode_rewards.append(eval_rewards)  # 添加奖励

            eval_rnn_states[eval_dones == True] = np.zeros(
                ((eval_dones == True).sum(), self.recurrent_N, self.hidden_size),
                dtype=np.float32,
            )  # 重置已完成环境的RNN状态
            eval_masks = np.ones((self.n_eval_rollout_threads, self.num_agents, 1), dtype=np.float32)  # 更新评估掩码
            eval_masks[eval_dones == True] = np.zeros(((eval_dones == True).sum(), 1), dtype=np.float32)  # 设置已完成环境的掩码为0

        eval_episode_rewards = np.array(eval_episode_rewards)  # 转换为NumPy数组
        eval_env_infos = {}  # 评估环境信息
        eval_env_infos["eval_average_episode_rewards"] = np.sum(np.array(eval_episode_rewards), axis=0)  # 计算总奖励
        eval_average_episode_rewards = np.mean(eval_env_infos["eval_average_episode_rewards"])  # 计算平均回合奖励
        print("智能体的评估平均回合奖励为: " + str(eval_average_episode_rewards))  # 打印平均回合奖励
        self.log_env(eval_env_infos, total_num_steps)  # 记录环境信息

    @torch.no_grad()
    def render(self):
        """
        可视化环境
        """
        envs = self.envs  # 获取环境

        all_frames = []  # 所有帧
        for episode in range(self.all_args.render_episodes):  # 遍历渲染回合
            obs = envs.reset()  # 重置环境
            if self.all_args.save_gifs:  # 如果保存GIF
                image = envs.render("rgb_array")[0][0]  # 渲染图像
                all_frames.append(image)  # 添加帧
            else:
                envs.render("human")  # 人类渲染

            rnn_states = np.zeros(
                (
                    self.n_rollout_threads,
                    self.num_agents,
                    self.recurrent_N,
                    self.hidden_size,
                ),
                dtype=np.float32,
            )  # 初始化RNN状态
            masks = np.ones((self.n_rollout_threads, self.num_agents, 1), dtype=np.float32)  # 初始化掩码

            episode_rewards = []  # 回合奖励

            for step in range(self.episode_length):  # 遍历步数
                calc_start = time.time()  # 记录开始时间

                self.trainer.prep_rollout()  # 准备收集数据
                action, rnn_states = self.trainer.policy.act(
                    np.concatenate(obs),  # 连接观察
                    np.concatenate(rnn_states),  # 连接RNN状态
                    np.concatenate(masks),  # 连接掩码
                    deterministic=True,  # 确定性动作
                )  # 获取动作
                actions = np.array(np.split(_t2n(action), self.n_rollout_threads))  # 分割动作
                rnn_states = np.array(np.split(_t2n(rnn_states), self.n_rollout_threads))  # 分割RNN状态

                if envs.action_space[0].__class__.__name__ == "MultiDiscrete":  # 如果是多离散动作空间
                    for i in range(envs.action_space[0].shape):
                        uc_actions_env = np.eye(envs.action_space[0].high[i] + 1)[actions[:, :, i]]  # one-hot编码
                        if i == 0:
                            actions_env = uc_actions_env  # 第一个动作
                        else:
                            actions_env = np.concatenate((actions_env, uc_actions_env), axis=2)  # 连接动作
                elif envs.action_space[0].__class__.__name__ == "Discrete":  # 如果是离散动作空间
                    actions_env = np.squeeze(np.eye(envs.action_space[0].n)[actions], 2)  # one-hot编码
                else:
                    raise NotImplementedError  # 不支持的动作空间

                # 观察奖励和下一个状态
                obs, rewards, dones, infos = envs.step(actions_env)  # 执行环境步进
                episode_rewards.append(rewards)  # 添加奖励

                rnn_states[dones == True] = np.zeros(
                    ((dones == True).sum(), self.recurrent_N, self.hidden_size),
                    dtype=np.float32,
                )  # 重置已完成环境的RNN状态
                masks = np.ones((self.n_rollout_threads, self.num_agents, 1), dtype=np.float32)  # 更新掩码
                masks[dones == True] = np.zeros(((dones == True).sum(), 1), dtype=np.float32)  # 设置已完成环境的掩码为0

                if self.all_args.save_gifs:  # 如果保存GIF
                    image = envs.render("rgb_array")[0][0]  # 渲染图像
                    all_frames.append(image)  # 添加帧
                    calc_end = time.time()  # 记录结束时间
                    elapsed = calc_end - calc_start  # 计算经过时间
                    if elapsed < self.all_args.ifi:  # 如果经过时间小于帧间隔
                        time.sleep(self.all_args.ifi - elapsed)  # 等待
                else:
                    envs.render("human")  # 人类渲染

            print("平均回合奖励为: " + str(np.mean(np.sum(np.array(episode_rewards), axis=0))))  # 打印平均回合奖励

        # if self.all_args.save_gifs:
        #     imageio.mimsave(str(self.gif_dir) + '/render.gif', all_frames, duration=self.all_args.ifi)
